# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/6 10:49 
@Author  : ZhangShenao 
@File    : 4.使用少样本提示模板.py 
@Desc    : 使用FewShotPromptTemplate少样本提示模板

Few-Shot Learning 少样本学习,是一种简单但强大的指导生成的方式
它可以降低Prompt的复杂度,通过少量的示例快速告知LLM的生成规范,显著提高模型性能
与Few-Shot对应的是Zero-Shot零样本
"""
import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate

dotenv.load_dotenv()

# 构造示例
examples = [
    {"question": "中国", "answer": "的首都是北京"},
    {"question": "美国", "answer": "的首都是华盛顿"},
    {"question": "日本", "answer": "的首都是东京"},
]

# 构造示例模板
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{question}"),
    ("ai", "{answer}"),
])

# 通过FewShotPromptTemplate少样本提示模板
few_shot_template = FewShotChatMessagePromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
)

# 构造最终提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个功能强大的聊天机器人，请按照严格示例提供的格式，回答用户的问题。"),
    few_shot_template,
    ("human", "{question}")
])

# 构造链并调用
llm = ChatZhipuAI(model="glm-4-flash")
chain = prompt | llm | StrOutputParser()

print(chain.invoke({"question": "英国"}))
